#!/usr/bin/env python2.7

# Copyright (c) 2016 Presto Labs Pte. Ltd.
# Author: jaewon

import argparse
import ctypes
import fcntl
import hashlib
import os
import pipes
import re
import signal
import subprocess
import sys
import tempfile
import time
import traceback

from datetime import datetime, timedelta
from distutils.spawn import find_executable  # pylint: disable=no-name-in-module


class InvalidFormatError(Exception):
  pass

class UserInterrupt(Exception):
  pass


def get_rsync_path():
  default_path = '/usr/bin/rsync'
  try:
    path = find_executable('rsync')
    return path or default_path
  except:  # pylint: disable=bare-except
    return default_path


def get_ssh_path():
  default_path = '/usr/bin/ssh'
  try:
    path = find_executable('ssh')
    return path or default_path
  except:  # pylint: disable=bare-except
    return default_path


def print_error(message):
  if sys.stderr.isatty():
    color_bold = '\033[1m'
    color_endc = '\033[0m'
    sys.stderr.write(color_bold + message + color_endc + '\n')
  else:
    sys.stderr.write(message + '\n')


# pylint: disable=C0122
def hhmm_to_second(hh_str, mm_str):
  hour = int(hh_str)
  minute = int(mm_str or '0')
  try:
    assert hour in range(0, 24+1)
    assert minute in range(0, 60)
    if hour == 24:
      assert minute == 0
  except AssertionError:
    raise InvalidFormatError()
  return hour * 3600 + minute * 60


def second_to_hhmm(second, sign=False):
  sign_letter = '-' if second < 0 else '+'
  if second >= 0 and not sign:
    sign_letter = ''
  second = abs(second)
  minute = (second % 3600) / 60
  hour = second / 3600
  return '%s%02d:%02d' % (sign_letter, hour, minute)


def parse_pause_period(string, tz_offset=0):
  """Returns list of intervals. Each interval consists of begin second and end
     second from midnight."""
  try:
    if string is None:
      return []
    regex_hhmm = '([0-9]?[0-9])(:([0-9][0-9]))?'
    regex_pause_time = '^%s-%s$' % (regex_hhmm, regex_hhmm)
    matched = re.match(regex_pause_time, string)
    assert matched is not None
    begin = ((hhmm_to_second(matched.group(1), matched.group(3))
              - tz_offset + 86400) % 86400)
    end = ((hhmm_to_second(matched.group(4), matched.group(6))
            - tz_offset + 86400) % 86400)
    if begin == end:
      return []
    elif begin < end:
      return [(begin, end)]
    elif begin > end:
      return [(begin, 86400)] + ([(0, end)] if end else [])
  except (AssertionError, InvalidFormatError):
    raise InvalidFormatError('fatal: Invalid pause period: %s' % string)


def parse_tz_offset(string):
  try:
    if string is None:
      return -time.timezone
    regex_hhmm = '([0-9]?[0-9])(:([0-9][0-9]))?'
    regex_tz_offset = r'^(\+|-)?%s$' % regex_hhmm
    matched = re.match(regex_tz_offset, string)
    assert matched is not None
    sign = matched.group(1) or '+'
    offset = hhmm_to_second(matched.group(2), matched.group(4))
    if sign == '+':
      return offset
    else:
      return -offset
  except (AssertionError, InvalidFormatError):
    raise InvalidFormatError('fatal: Invalid tz offset: %s' % string)


def generate_lock_id(src, dest, other=None):
  inputs = [src, dest]
  if other is not None:
    inputs.append(other)
  return hashlib.sha1('\t'.join(inputs)).hexdigest()


def parse_args():
  parser = argparse.ArgumentParser(
      description='A file copying tool. A wrapper of rsync but provides '
                  'additional features such as timeout, retry, pause-period '
                  'and lock based on (src, dest) string.')
  parser.add_argument('src', nargs=1, type=str, help='Source')
  parser.add_argument('dest', nargs=1, type=str, help='Destination')
  parser.add_argument('-q', '--quiet',
                      action='store_true', default=False,
                      help='Suppress non-error messages.')
  parser.add_argument('-v', '--verbose',
                      action='count', default=0,
                      help='Increase verbosity.')
  parser.add_argument('-a', '--append-verify',
                      action='store_true', default=False,
                      help='If specified, it will append data onto shorter '
                           'files with checksum.')
  parser.add_argument('-c', '--checksum',
                      action='store_true', default=False,
                      help='Skip based on checksum, not mod-time & size.')
  parser.add_argument('-b', '--batch-mode',
                      action='store_true', default=False,
                      help='If specified, it will not get user keyboard input '
                           'such as password. Useful when running in a script.')
  parser.add_argument('-m', '--max-retry',
                      action='store', type=int, default=5, metavar='N',
                      help='Maximum number of retry. If -1 is given, it will '
                           'retry until it succeeds.')
  parser.add_argument('--bwlimit',
                      action='store', type=int, metavar='KBytes/s',
                      help='limit I/O bandwidth; KBytes per second.')
  parser.add_argument('-t', '--timeout',
                      action='store', type=int,
                      help='Timeout in second. If pause time is given, '
                           'pause time is also regarded as running time.')
  parser.add_argument('--no-compression',
                      action='store_false', default=True, dest='compression',
                      help='If specified, it will not compress files when '
                           'filesync sends/receives the files.')
  parser.add_argument('-l', '--lock-id',
                      action='store', type=str,
                      help='If specified, it will wait if another remotecopy '
                           'with the same lock id is running. If not given, '
                           'it will generate lock id with src and dest '
                           'strings.')
  parser.add_argument('-p', '--pause-period',
                      action='append', type=str, metavar='HH:MM-HH:MM',
                      help='Specify pause periods in HH:MM-HH:MM format. '
                           'For example, if 08:00-16:00 is specified, it will '
                           'pause downloading files from 8AM to 4PM.')
  parser.add_argument('--tz-offset',
                      action='store', type=str, metavar='+HH:MM or -HH:MM',
                      help='Specify a timezone offset of the pause time in '
                           '+HH:MM or -HH:MM format.')
  parser.add_argument('-F', '--ssh-config',
                      action='store', type=str, metavar='SSH_CONFIG_FILE',
                      help='SSH configuration file.')
  parser.add_argument('-i', '--ssh-identity',
                      action='store', type=str, metavar='SSH_IDENTITY_FILE',
                      help='SSH identity file.')
  parser.add_argument('--ssh-strict-host-key-checking',
                      action='store_true', default=False,
                      help='If specified, it will do strict host key checking')
  return parser.parse_args()


# pylint: disable=R0902,R0912,R0913,R0914
def build_rsync_args(
    src,
    dest,
    append_verify=False,
    skip_on_checksum=False,
    batch_mode=False,
    bwlimit=None,
    compression=True,
    quiet=False,
    verbose_level=0,
    ssh_config_file=None,
    ssh_identity_file=None,
    ssh_strict_host_key_checking=False,
    rsync_path=None):
  rsync_args = [
      rsync_path or get_rsync_path(),
      '-a',  # --archive
      '-k'   # --copy-dirlinks
  ]

  if append_verify:
    rsync_args.append('--append-verify')

  if skip_on_checksum:
    rsync_args.append('--checksum')

  if bwlimit is not None:
    rsync_args.append('--bwlimit=%d' % bwlimit)

  if compression:
    rsync_args.append('-z')  # --compress

  if quiet:
    rsync_args.append('-q')  # --quiet
  else:
    rsync_args += [
        '-h',  # --human-readable
        '--progress'
    ]

  if verbose_level >= 1:
    rsync_args.append('-' + ('v' * verbose_level))

  # Build SSH args.
  ssh_args = []
  if ssh_config_file:
    ssh_args += ['-F', ssh_config_file]

  if ssh_identity_file:
    ssh_args += ['-i', ssh_identity_file]

  if batch_mode:
    ssh_args += ['-o', 'BatchMode=yes']

  if ssh_strict_host_key_checking:
    ssh_args += ['-o', 'StrictHostKeyChecking=yes']
  else:
    ssh_args += ['-o', 'StrictHostKeyChecking=no',
                 '-o', 'UserKnownHostsFile=/dev/null']

  if ssh_args:
    ssh_args = [get_ssh_path(), '-q'] + ssh_args
    ssh_cmd = ' '.join([pipes.quote(x) for x in ssh_args])
    rsync_args += ['-e', ssh_cmd]

  rsync_args += [src, dest]
  return rsync_args


# pylint: disable=W0702
def _set_death_signal():
  # Since prctl package is not system package we use in this way.
  # http://stackoverflow.com/questions/1884941/killing-the-children-with-the-parent
  libc_candidates = [
      '/lib64/libc.so.6',  # CentOS
      '/lib/x86_64-linux-gnu/libc.so.6'  # Ubuntu/Debian
  ]
  for libc_path in libc_candidates:
    try:
      libc = ctypes.CDLL(libc_path)
      libc.prctl(1, signal.SIGTERM)
      return
    except:  # Catches any exception.
      # If it cannot find relevant libraris, just skip.
      continue


class RsyncController(object):
  def __init__(self, cli_args):
    self._rsync_args = build_rsync_args(
        src=cli_args.src[0],
        dest=cli_args.dest[0],
        append_verify=cli_args.append_verify,
        skip_on_checksum=cli_args.checksum,
        batch_mode=cli_args.batch_mode,
        bwlimit=cli_args.bwlimit,
        compression=cli_args.compression,
        quiet=cli_args.quiet,
        verbose_level=cli_args.verbose,
        ssh_config_file=cli_args.ssh_config,
        ssh_identity_file=cli_args.ssh_identity,
        ssh_strict_host_key_checking=cli_args.ssh_strict_host_key_checking)
    self._popen = None

  def is_running(self):
    if self._popen is not None:
      self._popen.poll()
      return self._popen.returncode is None
    return False

  def send_signal(self, signum):
    if self._popen is not None:
      self._popen.send_signal(signum)

  @property
  def command(self):
    return ' '.join([pipes.quote(x) for x in self._rsync_args])

  @property
  def returncode(self):
    assert self._popen is not None
    return self._popen.returncode

  def start(self):
    assert not self.is_running()
    try:
      self._popen = subprocess.Popen(self._rsync_args,
                                     preexec_fn=_set_death_signal)
    except Exception as e:
      print_error('error: %s' % e.message)
      raise

  def stop(self, kill_timeout=3, send_sigterm=True):
    if not self.is_running():
      return
    if send_sigterm:
      self._popen.terminate()
    until = datetime.now() + timedelta(seconds=kill_timeout)
    while self.is_running() and datetime.now() < until:
      time.sleep(0.1)
    if self.is_running():
      try:
        self._popen.terminate()
        self._popen.wait()
      except OSError as e:
        print_error('warning: %s' % e.message)
      except Exception as e:
        print_error('error: %s' % e.message)
        raise


class RsyncExecutor(object):
  STATE_INITIAL = 1
  STATE_RUNNING = 2
  STATE_PAUSED = 3
  STATE_RETRY = 4
  STATE_TIMEOUT = 5

  LOCKFILE_DIR = os.path.join(tempfile.gettempdir(), 'remotecopy')

  def __init__(self, cli_args):
    self._controller = RsyncController(cli_args)
    self._exitsignal = False
    self._exitcode = 0
    self._state = self.STATE_INITIAL
    self._retry_count = 0
    self._max_retry = cli_args.max_retry

    self._timeout = cli_args.timeout
    self._run_until = None

    self._lock_id = (cli_args.lock_id or
                     generate_lock_id(cli_args.src[0], cli_args.dest[0]))
    self._lockfile = None

    tz_offset = parse_tz_offset(cli_args.tz_offset)
    self._pause_period = []
    if cli_args.pause_period:
      for period in cli_args.pause_period:
        self._pause_period += parse_pause_period(period, tz_offset)

    self._dispatch = {
        self.STATE_INITIAL: self._handle_not_running,
        self.STATE_PAUSED: self._handle_not_running,
        self.STATE_RETRY: self._handle_not_running,
        self.STATE_RUNNING: self._handle_running,
        self.STATE_TIMEOUT: self._handle_timeout
    }

    if cli_args.verbose >= 1:
      sys.stdout.write('rsync command: %s\n' % self._controller.command)
      if self._pause_period:
        sys.stdout.write('pause period:\n')
        for interval \
            in sum([parse_pause_period(p) for p in cli_args.pause_period], []):
          tz_str = ''
          if cli_args.tz_offset is not None:
            tz_str = ' %s' % second_to_hhmm(tz_offset, True)
          sys.stdout.write(
              '  %s-%s%s\n' % (second_to_hhmm(interval[0]),
                               second_to_hhmm(interval[1]),
                               tz_str))
      sys.stdout.write('\n')

  def _open_lockfile(self):
    if not os.path.isdir(self.LOCKFILE_DIR):
      try:
        os.makedirs(self.LOCKFILE_DIR, 0777)
        os.chmod(self.LOCKFILE_DIR, 0777)
      except OSError:
        pass
    assert os.path.isdir(self.LOCKFILE_DIR), "fatal: cannot create lock file."

    path = os.path.join(self.LOCKFILE_DIR, self._lock_id)
    if not os.path.isfile(path):
      try:
        prev_umask = os.umask(0000)
        flags = os.O_CREAT | os.O_APPEND
        lock_fd = os.open(path, flags, 0666)
        os.umask(prev_umask)
        os.close(lock_fd)
      except OSError:
        pass
    self._lockfile = open(path, 'w+')

  def _try_lock(self):
    try:
      fcntl.lockf(self._lockfile, fcntl.LOCK_EX | fcntl.LOCK_NB)
      return True
    except IOError:
      return False

  def _unlock(self):
    fcntl.lockf(self._lockfile, fcntl.LOCK_UN)

  def is_pause_time(self):
    utc_dt = datetime.utcnow()
    dt_sec = utc_dt.hour * 3600 + utc_dt.minute * 60
    for period in self._pause_period:
      if dt_sec in range(period[0], period[1]):
        return True
    return False

  def _set_timeout(self):
    if self._timeout is None:
      self._run_until = datetime.max
    else:
      self._run_until = datetime.now() + timedelta(seconds=self._timeout)

  def _check_timeout(self):
    if datetime.now() >= self._run_until:
      self._state = self.STATE_TIMEOUT

  def _exit(self, exitcode):
    self._exitsignal = True
    self._exitcode = exitcode

  def _handle_not_running(self):
    if not self.is_pause_time() and self._try_lock():
      try:
        if self._state == self.STATE_PAUSED:
          print_error('Resuming...')
        elif self._state == self.STATE_RETRY:
          if self._max_retry != -1:
            print_error('Retrying (%d/%d)...'
                        % (self._retry_count, self._max_retry))
          else:
            print_error('Retrying (%d)...'  % self._retry_count)
        self._controller.start()
        self._state = self.STATE_RUNNING
      except OSError:
        self._unlock()
        self._exit(2)

  def _handle_running(self):
    if not self._controller.is_running():
      self._unlock()
      if self._controller.returncode == 0:
        self._exit(0)
        return

      self._retry_count += 1
      if self._retry_count <= self._max_retry or self._max_retry == -1:
        print_error('rsync is dead unexpectedly (%d). '
                    'Retrying in 5 seconds...'
                    % self._controller.returncode)
        self._state = self.STATE_RETRY
        time.sleep(4)
      else:
        print_error('rsync is dead unexpectedly (%d). Quiting...'
                    % self._controller.returncode)
        self._exit(1)
      return

    if self.is_pause_time():
      self._controller.stop()
      self._unlock()
      self._state = self.STATE_PAUSED
      print_error('Paused.')

  def _handle_timeout(self):
    if self._controller.is_running():
      self._controller.stop()
      self._unlock()
    print_error('Timeout. terminating...')
    self._exit(3)

  # pylint: disable=W0613
  def _handle_signal(self, signum, frame):
    if self._controller.is_running():
      self._controller.send_signal(signum)
      self._controller.stop(send_sigterm=False)
    print_error('Signaled. terminating...')
    raise UserInterrupt()

  def run(self):
    self._set_timeout()
    signal.signal(signal.SIGINT, self._handle_signal)
    signal.signal(signal.SIGTERM, self._handle_signal)
    signal.signal(signal.SIGHUP, self._handle_signal)

    try:
      self._open_lockfile()
      while True:
        self._check_timeout()
        self._dispatch[self._state]()
        if self._exitsignal:
          break
        time.sleep(1)
      return self._exitcode
    except UserInterrupt:
      return 1
    finally:
      self._lockfile.close()


# pylint: disable=W0703
def main():
  args = parse_args()
  try:
    executor = RsyncExecutor(args)
    return executor.run()

  except Exception as e:
    sys.stderr.write('%s\n' % e.message)
    if args.verbose > 0:
      sys.stderr.write('\n---\n')
      traceback.print_exc(file=sys.stderr)
    return 9


if __name__ == '__main__':
  sys.exit(main())
